#pragma once
#ifndef DWCONV_H_
#define DWCONV_H_

#include "stream_tools.h"

// depthwise conv3x3
template<
    unsigned N_IO,     // # of IO
    unsigned N_CH,     // # of CHannel
    unsigned BIT_IN,    // nBit of INput
    unsigned BIT_WT,    // nBit of WeighT
    unsigned BIT_OUT,   // nBit of OUTput
    unsigned VEC_LEN    // ROW * COL
>
void dwconv_3x3(data_stream<N_IO * BIT_IN>& in, data_stream<N_IO * BIT_OUT>& out, const ap_uint<N_IO * BIT_WT> weight[N_CH / N_IO][9])
// origin:    ap_int<BIT_WT> wt[9][N_CH]
// reshape:   ap_int<BIT_WT> wt[9][N_CH / N_IO][N_IO]
// transpose: ap_int<BIT_WT> wt[N_CH / N_IO][9][N_IO]
// reshape:   ap_uint<N_IO * BIT_WT> wt[N_CH / N_IO][9]
{
    static_assert(N_CH >= N_IO, "dwconv_3x3");
    static_assert(N_CH % N_IO == 0, "dwconv_3x3");
    constexpr unsigned FOLD = N_CH / N_IO;
    constexpr unsigned ITERS = VEC_LEN;

    assert(in.size() == VEC_LEN * FOLD * 9);
    assert(out.empty());

    #pragma HLS DATAFLOW

    ap_int<BIT_OUT> acc[N_IO];
    #pragma HLS ARRAY_PARTITION variable=acc complete dim=1

    for (unsigned i = 0; i < N_IO; ++i)
    {
        #pragma HLS UNROLL
        acc[i] = 0;
    }

    for (unsigned it = 0; it < ITERS; ++it)
    {
        for (unsigned f = 0; f < FOLD; ++f)
        {
            for (unsigned k = 0; k < 9; ++k)
            {
                #pragma HLS PIPELINE II=1
                // load
                ap_uint<N_IO * BIT_IN> in_buf = in.read();
                ap_uint<N_IO * BIT_WT> wt_buf = weight[f][k];

                // calc
                for (unsigned i = 0; i < N_IO; ++i)
                {
                    #pragma HLS UNROLL
                    ap_uint<BIT_IN> x = in_buf(SLICE(BIT_IN, i));
                    ap_int<BIT_WT>  y = wt_buf(SLICE(BIT_WT, i));
                    acc[i] += x * y;
                }

                // output
                if (k == 8)
                {
                    ap_uint<N_IO * BIT_OUT> out_buf;
                    for (unsigned i = 0; i < N_IO; ++i)
                    {
                        #pragma HLS UNROLL
                        out_buf(SLICE(BIT_OUT, i)) = acc[i];
                        acc[i] = 0;
                    }
                    out.write(out_buf);
                }
            }
        }
    }

    assert(in.empty());
    assert(out.size() == VEC_LEN * FOLD);
    return;
}

template<
    unsigned W_IO,  // Width of IO      N_IO * BIT
    unsigned W_CH,  // Width of CHannel N_CH * BIT
    unsigned ROW,
    unsigned COL,
    unsigned N_BATCH,
    unsigned PADVAL = 0x00
>
void im2col_3x3(data_stream<W_IO>& in, data_stream<W_IO>& out)
{
    static_assert(W_CH >= W_IO, "im2col");
    static_assert(W_CH % W_IO == 0, "im2col");
    constexpr unsigned FOLD = W_CH / W_IO;
    constexpr unsigned COL_OFFSET = 2;
    const ap_uint<W_IO> PAD_AP = PADVAL;

    assert(in.size() == N_BATCH * ROW * COL * FOLD);
    assert(out.empty());

    #pragma HLS DATAFLOW

    static ap_uint<2> idx[3] = {0, 1, 2};
    #pragma HLS ARRAY_PARTITION variable=idx complete dim=1
    ap_uint<W_IO> line[3][COL][FOLD];
    
    for (unsigned c = 0; c < COL + COL_OFFSET; ++c)
    {
        for (unsigned f = 0; f < FOLD; ++f)
        {
            #pragma HLS PIPELINE II=1
            if (c < COL)
            {
                line[idx[2]][c][f] = in.read();
            }
            else
            {
                line[idx[0]][c - COL][f] = in.read();
            }
        }
    }

    // main loop
    for (unsigned b = 0; b < N_BATCH; ++b)
    {
        for (unsigned r = 0; r < ROW; ++r)
        {
            for (unsigned c = 0; c < COL; ++c)
            {
                for (unsigned f = 0; f < FOLD; ++f)
                {
                    #pragma HLS PIPELINE II=9
                    // idx shift
                    if (c == 0 && f == 0)
                    {
                        ap_uint<2> it = idx[0];
                        idx[0] = idx[1];
                        idx[1] = idx[2];
                        idx[2] = it;
                    }
                    // load data
                    if (b < N_BATCH - 1 || r < ROW - 2 || (r == ROW - 2 && c < COL - COL_OFFSET))
                    {
                        ap_uint<W_IO> in_buf = in.read();
                        if (c < COL - COL_OFFSET)
                        {
                            line[idx[2]][c + COL_OFFSET][f] = in_buf;
                        }
                        else
                        {
                            line[idx[0]][c + COL_OFFSET - COL][f] = in_buf;
                        }
                    }

                    // output, II = 9
                    // (r == 0 || r == ROW - 1 || c == 0 || c == COL - 1) --> padding
                    out.write((r != 0 && c != 0) ? 
                        line[idx[0]][c - 1][f] : PAD_AP);
                    out.write((r != 0) ? 
                        line[idx[0]][c + 0][f] : PAD_AP);
                    out.write((r != 0 && c != COL - 1)? 
                        line[idx[0]][c + 1][f] : PAD_AP);
                    out.write((c != 0) ? 
                        line[idx[1]][c - 1][f] : PAD_AP);
                    out.write(
                        line[idx[1]][c + 0][f]);
                    out.write((c != COL - 1) ? 
                        line[idx[1]][c + 1][f] : PAD_AP);
                    out.write((r != ROW - 1 && c != 0) ? 
                        line[idx[2]][c - 1][f] : PAD_AP);
                    out.write((r != ROW - 1) ? 
                        line[idx[2]][c + 0][f] : PAD_AP);
                    out.write((r != ROW - 1 && c != COL - 1) ? 
                        line[idx[2]][c + 1][f] : PAD_AP);
                }
            }
        }
    }

    assert(in.empty());
    assert(out.size() == N_BATCH * ROW * COL * FOLD * 9);
    return;
}

// depthwise conv3x3
template<
    unsigned N_IO,     // # of IO
    unsigned N_CH,     // # of CHannel
    unsigned BIT_IN,    // nBit of INput
    unsigned BIT_WT,    // nBit of WeighT
    unsigned BIT_OUT,   // nBit of OUTput
    unsigned VEC_LEN    // ROW * COL
>
void dwconv_3x3_old(data_stream<N_IO * BIT_IN>& in, data_stream<N_IO * BIT_OUT>& out, const ap_uint<N_IO * BIT_WT> weight[N_CH / N_IO][9])
// origin:    ap_int<BIT_WT> wt[9][N_CH]
// reshape:   ap_int<BIT_WT> wt[9][N_CH / N_IO][N_IO]
// transpose: ap_int<BIT_WT> wt[N_CH / N_IO][9][N_IO]
// reshape:   ap_uint<N_IO * BIT_WT> wt[N_CH / N_IO][9]
{
    static_assert(N_CH >= N_IO, "dwconv_3x3");
    static_assert(N_CH % N_IO == 0, "dwconv_3x3");
    constexpr unsigned FOLD = N_CH / N_IO;
    constexpr unsigned ITERS = VEC_LEN * FOLD * 9;

    assert(in.size() == VEC_LEN * FOLD * 9);
    assert(out.empty());

    unsigned f = 0;
    unsigned k = 0;
    ap_int<BIT_OUT> acc[N_IO];
#pragma HLS ARRAY_PARTITION variable=acc complete dim=1
    for (unsigned i = 0; i < N_IO; ++i)
    {
#pragma HLS UNROLL
        // UnRoll
        acc[i] = 0;
    }

    for (unsigned it = 0; it < ITERS; ++it)
    {
#pragma HLS PIPELINE II=1
        // II = 1
        ap_uint<N_IO * BIT_IN> in_buf = in.read();
        ap_uint<N_IO * BIT_WT> wt_buf = weight[f][k];
        for (unsigned i = 0; i < N_IO; ++i)
        {
#pragma HLS UNROLL
            // UnRoll
            ap_uint<BIT_IN> x = in_buf(SLICE(BIT_IN, i));
            ap_int<BIT_WT>  y = wt_buf(SLICE(BIT_WT, i));
            acc[i] += x * y;
        }
        if (k == 8)
        {
            ap_uint<N_IO * BIT_OUT> out_buf;
            for (unsigned i = 0; i < N_IO; ++i)
            {
#pragma HLS UNROLL
                // UnRoll
                out_buf(SLICE(BIT_OUT, i)) = acc[i];
                acc[i] = 0;
            }
            out.write(out_buf);
            k = 0;
            f = (f != FOLD - 1) ? f + 1 : 0;
        }
        else
        {
            ++k;
        }
    }

    assert(in.empty());
    assert(out.size() == VEC_LEN * FOLD);
    return;
}

#endif
